# -*- coding: utf-8 -*-
# @Author: lidongdong
# @time  : 19-1-23 下午10:00
# @file  : utils.py

import os
import json
from plainobj import CaptionUnit
from tqdm import tqdm
import numpy as np
import h5py
import cv2
"""
统计下列形式的json信息
{
    "h5py_file": "val.h5",
    "mscoco_val":
    [{
    "index": 0,
    "file_name": "COCO_val2014_000000391895.jpg",
    "captions": ["A black Honda motorcycle parked in front of a garage.",
                "A view of train tracks and a train."],
    "caption_split":[[], [], [], [], []]
    },
    {
    "id": 2,
    ...
    }]
]
"""


def prepare_new_json(json_path, mode, new_json_path=None):
    with open(json_path) as f:
        origin = json.load(f)
    images = origin["images"]
    annotations =origin["annotations"]
    images = sorted(images, key=lambda x: x["id"])
    image_annotations_dict = {}
    for image in images:
        image_annotations_dict[image["id"]] = []
    for annotation in annotations:
        image_id = annotation["image_id"]
        image_annotations_dict[image_id].append(annotation)

    # body
    caption_body = []
    for i, key in tqdm(enumerate(sorted(image_annotations_dict.keys()))):
        image = images[i]
        # prepare CapitionUnit
        anns = image_annotations_dict[key]
        index = i
        filename = image["file_name"]
        captions = [ann["caption"] for ann in anns]
        caption_body.append(CaptionUnit(index, filename, captions).__dict__)

    json_content = {}
    json_content["h5py_file"] = "{}.h5".format(mode)
    json_content["content"] = caption_body

    if new_json_path is not None:
        with open(new_json_path, "w") as f:
            json.dump(json_content, f, indent=4)

    return json_content, caption_body


def image2h5(image_filenames, h5filename, root=None):
    """
    load all image to h5py file
    :param image_filenames: filenames of all image
    :param h5filename: the h5py file name
    :param root: the root path of images
    :return:
    """
    # TODO 将image filenames 指定的image保存到h5filename中 dubug!!!
    h5file = h5py.File(h5filename, "w")
    image_dataset = h5file.create_dataset("images",
                                          shape=(0, 256, 256, 3),
                                          maxshape=(None, 256, 256, 3),
                                          compression="gzip",
                                          chunks=(1, 64, 64, 3),
                                          dtype=np.float32)

    if root is not None:
        image_filenames = map(lambda x: os.path.join(root, x), image_filenames)
    print "loading data, write to h5file"
    step = 100
    temp = []
    for i in range(0, len(image_filenames), step):
        temp.append(image_filenames[i: i + step])

    image_filenames = temp
    for filenames in tqdm(image_filenames):
        image = load_batch_images(filenames)
        new_batch_size = image.shape[0]
        if new_batch_size == 0:
            break
        original_size = image_dataset.shape[0]
        image_dataset.resize(original_size + new_batch_size, axis=0)
        image_dataset[original_size: original_size + new_batch_size, :, :, :] = image

    h5file.close()
    print "Done"


def load_one_image(filename):
    image = cv2.imread(filename)
    image = cv2.resize(image, (256, 256))
    image = image.astype(np.uint8)
    image = image[:, :, ::-1]
    return image


def load_batch_images(filenames):
    images = []
    for filename in filenames:
        images.append(load_one_image(filename))
    return np.asarray(images)


if __name__ == '__main__':
    json_content, caption_body = prepare_new_json("/dl_data/mscoco/annotations_trainval2014/annotations/captions_val2014.json",
                     "val", "/dl_data/mscoco/annotations_trainval2014/annotations/val2014.json")
    h5py_file = json_content["h5py_file"]
    h5py_file = os.path.join("/dl_data/mscoco/", h5py_file)
    filenames = [caption["filename"] for caption in caption_body]
    val_root = "/dl_data/mscoco/val2014/"
    image_filenames = map(lambda x: os.path.join(val_root, x), filenames)
    image2h5(image_filenames, h5py_file)
    print "DONE"
